#include "data_structures.hpp"
#include <cmath>
#include <algorithm>
#include <iostream>

#ifndef M_PI
#define M_PI 3.14159265358979323846
#endif

#ifdef _OPENMP
#include <omp.h>
#endif

namespace EddyDetection {

namespace {
    const double R_EARTH = 6371e3;
    const double DEG_TO_RAD = M_PI / 180.0;
    const double RAD_TO_DEG = 180.0 / M_PI;

    double bilinearInterpolation(const Eigen::MatrixXd& field,
                                 const Eigen::VectorXd& x_coords,
                                 const Eigen::VectorXd& y_coords,
                                 double x, double y) {

        int i = 0, j = 0;
        for (i = 0; i < x_coords.size() - 1; ++i) {
            if (x >= x_coords[i] && x <= x_coords[i + 1]) break;
        }
        for (j = 0; j < y_coords.size() - 1; ++j) {
            if (y >= y_coords[j] && y <= y_coords[j + 1]) break;
        }
        
        if (i >= x_coords.size() - 1) i = x_coords.size() - 2;
        if (j >= y_coords.size() - 1) j = y_coords.size() - 2;
        
        double x1 = x_coords[i], x2 = x_coords[i + 1];
        double y1 = y_coords[j], y2 = y_coords[j + 1];
        
        double Q11 = field(i, j);
        double Q12 = field(i, j + 1);  
        double Q21 = field(i + 1, j);
        double Q22 = field(i + 1, j + 1);
        
        double wx = (x - x1) / (x2 - x1);
        double wy = (y - y1) / (y2 - y1);
        
        return Q11 * (1 - wx) * (1 - wy) + 
               Q21 * wx * (1 - wy) + 
               Q12 * (1 - wx) * wy + 
               Q22 * wx * wy;
    }

    Eigen::Vector2d rk4Step(const Eigen::Vector2d& pos, double dt,
                            const Eigen::MatrixXd& u_field, const Eigen::MatrixXd& v_field,
                            const Eigen::VectorXd& lat, const Eigen::VectorXd& lon) {
        
        auto getVelocity = [&](const Eigen::Vector2d& p) -> Eigen::Vector2d {
            double u_val = bilinearInterpolation(u_field, lat, lon, p[0], p[1]);
            double v_val = bilinearInterpolation(v_field, lat, lon, p[0], p[1]);
            
            double lat_rad = p[1] * DEG_TO_RAD;
            double deg_per_meter = RAD_TO_DEG / R_EARTH;
            
            double dlon_dt = (u_val / std::cos(lat_rad)) * deg_per_meter * 3600.0;
            double dlat_dt = v_val * deg_per_meter * 3600.0;
            
            return Eigen::Vector2d(dlon_dt, dlat_dt);
        };
        
        Eigen::Vector2d k1 = getVelocity(pos);
        Eigen::Vector2d k2 = getVelocity(pos + k1 * dt * 0.5);
        Eigen::Vector2d k3 = getVelocity(pos + k2 * dt * 0.5);
        Eigen::Vector2d k4 = getVelocity(pos + k3 * dt);
        
        return pos + (k1 + 2*k2 + 2*k3 + k4) * dt / 6.0;
    }

    double calculatePercentileThreshold(const Eigen::MatrixXd& field, double percentile) {
        std::vector<double> values;
        values.reserve(field.rows() * field.cols());
        
        for (int i = 0; i < field.rows(); ++i) {
            for (int j = 0; j < field.cols(); ++j) {
                if (std::isfinite(field(i, j))) {
                    values.push_back(field(i, j));
                }
            }
        }
        
        if (values.empty()) return 0.0;
        
        std::sort(values.begin(), values.end());
        size_t index = static_cast<size_t>(percentile / 100.0 * (values.size() - 1));
        return values[index];
    }
}

Eigen::MatrixXd LagrangianCore::computeFTLE(const VelocityField& velocity_field,
                                            const EddyDetectionConfig& config,
                                            int time_index) {
    const auto& lat = velocity_field.lat;
    const auto& lon = velocity_field.lon;
    const auto& u_field = velocity_field.u_layers[0];
    const auto& v_field = velocity_field.v_layers[0];
    
    if (u_field.size() == 0 || v_field.size() == 0) {
        std::cerr << "Error: Speed field data is empty" << std::endl;
        return Eigen::MatrixXd();
    }
    
    int nx = lat.size();
    int ny = lon.size();

    Eigen::MatrixXd initial_x(nx, ny);
    Eigen::MatrixXd initial_y(nx, ny);
    
    for (int i = 0; i < nx; ++i) {
        for (int j = 0; j < ny; ++j) {
            initial_x(i, j) = lat[i];
            initial_y(i, j) = lon[j];
        }
    }

    Eigen::MatrixXd final_x(nx, ny);
    Eigen::MatrixXd final_y(nx, ny);
    
    double dt = 0.1;
    int num_steps = static_cast<int>(config.integration_time_hours / dt);
    
#ifdef _OPENMP
    if (config.use_openmp) {
        std::cout << "duoxiancheng is open " << std::endl;
        int num_threads = config.num_threads > 0 ? config.num_threads : omp_get_max_threads();
        omp_set_num_threads(num_threads);
        
        #pragma omp parallel for collapse(2)
        for (int i = 0; i < nx; ++i) {
            auto startTime = std::chrono::high_resolution_clock::now();
            for (int j = 0; j < ny; ++j) {
                Eigen::Vector2d pos(initial_x(i, j), initial_y(i, j));

                for (int step = 0; step < num_steps; ++step) {
                    pos = rk4Step(pos, dt, u_field, v_field, lat, lon);
                }
                
                final_x(i, j) = pos[0];
                final_y(i, j) = pos[1];
            }
            auto endTime = std::chrono::high_resolution_clock::now();
            auto duraTime = std::chrono::duration_cast<std::chrono::milliseconds>(endTime - startTime);
            std::cout << "the time of nx[" << i << "] = " << duraTime.count() << "ms" << std::endl;
        }
    } else
#endif
    {
        for (int i = 0; i < nx; ++i) {
            for (int j = 0; j < ny; ++j) {
                Eigen::Vector2d pos(initial_x(i, j), initial_y(i, j));
                
                for (int step = 0; step < num_steps; ++step) {
                    pos = rk4Step(pos, dt, u_field, v_field, lat, lon);
                }
                
                final_x(i, j) = pos[0];
                final_y(i, j) = pos[1];
            }
        }
    }

    Eigen::MatrixXd ftle_field(nx, ny);
    
    double dx = (lat.size() > 1) ? (lat[1] - lat[0]) : 1.0;
    double dy = (lon.size() > 1) ? (lon[1] - lon[0]) : 1.0;
    
    for (int i = 1; i < nx - 1; ++i) {
        for (int j = 1; j < ny - 1; ++j) {

            double dphi_x_dx = (final_x(i+1, j) - final_x(i-1, j)) / (2.0 * dx);
            double dphi_x_dy = (final_x(i, j+1) - final_x(i, j-1)) / (2.0 * dy);
            double dphi_y_dx = (final_y(i+1, j) - final_y(i-1, j)) / (2.0 * dx);
            double dphi_y_dy = (final_y(i, j+1) - final_y(i, j-1)) / (2.0 * dy);
            
            double C11 = dphi_x_dx * dphi_x_dx + dphi_y_dx * dphi_y_dx;
            double C12 = dphi_x_dx * dphi_x_dy + dphi_y_dx * dphi_y_dy;
            double C22 = dphi_x_dy * dphi_x_dy + dphi_y_dy * dphi_y_dy;
            
            double trace = C11 + C22;
            double det = C11 * C22 - C12 * C12;
            double lambda_max = 0.5 * (trace + std::sqrt(std::max(0.0, trace*trace - 4*det)));

            ftle_field(i, j) = std::log(std::sqrt(std::max(1.0, lambda_max))) / config.integration_time_hours;
        }
    }

    for (int i = 0; i < nx; ++i) {
        ftle_field(i, 0) = 0.0;
        ftle_field(i, ny-1) = 0.0;
    }
    for (int j = 0; j < ny; ++j) {
        ftle_field(0, j) = 0.0;
        ftle_field(nx-1, j) = 0.0;
    }
    std::cout << "ftle_field.size = " << ftle_field.size() << std::endl;

    return ftle_field;
}


std::vector<std::vector<Eigen::Vector2d>> LagrangianCore::findFTLEContours(
    const Eigen::MatrixXd& ftle_field,
    const Eigen::VectorXd& lat, const Eigen::VectorXd& lon,
    double threshold_percentile) {

    double threshold = calculatePercentileThreshold(ftle_field, threshold_percentile);

    std::vector<std::vector<Eigen::Vector2d>> contours;

    Eigen::MatrixXi binary_image = Eigen::MatrixXi::Zero(ftle_field.rows(), ftle_field.cols());
    for (int i = 0; i < ftle_field.rows(); ++i) {
        for (int j = 0; j < ftle_field.cols(); ++j) {
            binary_image(i, j) = (ftle_field(i, j) >= threshold) ? 1 : 0;
        }
    }

    Eigen::MatrixXi visited = Eigen::MatrixXi::Zero(binary_image.rows(), binary_image.cols());
    
    for (int i = 1; i < binary_image.rows() - 1; ++i) {
        for (int j = 1; j < binary_image.cols() - 1; ++j) {
            if (binary_image(i, j) == 1 && visited(i, j) == 0) {
                std::vector<Eigen::Vector2d> contour;

                std::vector<std::pair<int, int>> region_pixels;
                std::vector<std::pair<int, int>> stack;
                stack.push_back({i, j});
                
                while (!stack.empty()) {
                    auto [x, y] = stack.back();
                    stack.pop_back();
                    
                    if (x < 0 || x >= binary_image.rows() || y < 0 || y >= binary_image.cols() ||
                        visited(x, y) == 1 || binary_image(x, y) == 0) {
                        continue;
                    }
                    
                    visited(x, y) = 1;
                    region_pixels.push_back({x, y});

                    stack.push_back({x+1, y});
                    stack.push_back({x-1, y});
                    stack.push_back({x, y+1});
                    stack.push_back({x, y-1});
                }

                for (const auto& pixel : region_pixels) {
                    if (pixel.first > 0 && pixel.first < binary_image.rows()-1 &&
                        pixel.second > 0 && pixel.second < binary_image.cols()-1) {

                        bool is_boundary = false;
                        for (int di = -1; di <= 1; ++di) {
                            for (int dj = -1; dj <= 1; ++dj) {
                                if (binary_image(pixel.first + di, pixel.second + dj) == 0) {
                                    is_boundary = true;
                                    break;
                                }
                            }
                            if (is_boundary) break;
                        }
                        
                        if (is_boundary) {
                            double x_coord = lat[pixel.first];
                            double y_coord = lon[pixel.second];
                            contour.push_back(Eigen::Vector2d(x_coord, y_coord));
                        }
                    }
                }
                
                if (contour.size() >= 4) {
                    contours.push_back(contour);
                }
            }
        }
    }
    
    return contours;
}

std::vector<EddyContour> LagrangianCore::detectEddiesLagrangian(
    const VelocityField& velocity_field,
    const TemperatureField& temp_field,
    const EddyDetectionConfig& config,
    int time_index) {
    
    std::cout << "Start Lagrange vortex detection..." << std::endl;

    Eigen::MatrixXd ftle_field = computeFTLE(velocity_field, config, time_index);

    auto contours = findFTLEContours(ftle_field, velocity_field.lat, velocity_field.lon,
                                     config.ftle_threshold_percentile);
    
    std::cout << "Find " << contours.size() << " FTLE contours" << std::endl;

    Eigen::MatrixXd vorticity = EulerianCore::computeVorticity(velocity_field, time_index);

    std::vector<EddyContour> eddies;
    for (const auto& contour : contours) {
        EddyContour eddy = EddyAnalyzer::analyzeEddy(contour, vorticity, temp_field.data_layers[0],
                                                     velocity_field.lat, velocity_field.lon,
                                                     config.min_area_km2);
        if (eddy.area_km2 >= config.min_area_km2) {
            eddies.push_back(eddy);
        }
    }

    eddies = EddyAnalyzer::filterAndSortEddies(eddies, config.min_area_km2, velocity_field);
    
    std::cout << "Detected " << eddies.size() << " valid vortices" << std::endl;
    return eddies;
}

Eigen::MatrixXd LagrangianCore::computeFlowMap(const VelocityField& velocity_field,
                                               double t_start, double t_end,
                                               const EddyDetectionConfig& config) {
    // 这个函数可以在未来扩展为更复杂的流映射计算
    // 目前由computeFTLE内部处理
    return Eigen::MatrixXd();
}

} // namespace EddyDetection 